{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "1d1fbb9e6c6649188b3eb8a7ff64ea40": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_3b7017e547ae4e418b6296f6421d8a36",
              "IPY_MODEL_62d29109e30148cd8d33b9d0b58d5b81",
              "IPY_MODEL_30694df5cbdb4e00b70534f963ec61fd"
            ],
            "layout": "IPY_MODEL_11ff540c531c48f08b241beb8e5e3a19"
          }
        },
        "3b7017e547ae4e418b6296f6421d8a36": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_36407cf7a73d466abbe732b7681a4de1",
            "placeholder": "​",
            "style": "IPY_MODEL_7bd1f4447764435cbbdcde2d587906cb",
            "value": "Batches: 100%"
          }
        },
        "62d29109e30148cd8d33b9d0b58d5b81": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_a517d1a8c2de4333850843670a8505ba",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_f4064c8c509e4bc0887a2222e9b2c3e2",
            "value": 1
          }
        },
        "30694df5cbdb4e00b70534f963ec61fd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7ce67680963c4c1fafc2df67bfe97710",
            "placeholder": "​",
            "style": "IPY_MODEL_e708fe1190174d7597082288d6a9cc1a",
            "value": " 1/1 [00:02&lt;00:00,  2.66s/it]"
          }
        },
        "11ff540c531c48f08b241beb8e5e3a19": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "36407cf7a73d466abbe732b7681a4de1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7bd1f4447764435cbbdcde2d587906cb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "a517d1a8c2de4333850843670a8505ba": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "f4064c8c509e4bc0887a2222e9b2c3e2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7ce67680963c4c1fafc2df67bfe97710": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "e708fe1190174d7597082288d6a9cc1a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uVehGld9c3mq",
        "outputId": "57552217-3102-47f5-de2a-a1254d6d80f9"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/67.3 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m67.3/67.3 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m88.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.0/1.0 MB\u001b[0m \u001b[31m54.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.2/1.2 MB\u001b[0m \u001b[31m61.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m615.7/615.7 kB\u001b[0m \u001b[31m42.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.4/2.4 MB\u001b[0m \u001b[31m89.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m298.0/298.0 kB\u001b[0m \u001b[31m26.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m273.8/273.8 kB\u001b[0m \u001b[31m25.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m94.7/94.7 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m69.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m408.7/408.7 kB\u001b[0m \u001b[31m34.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m93.2/93.2 kB\u001b[0m \u001b[31m8.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.3/13.3 MB\u001b[0m \u001b[31m36.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m52.5/52.5 kB\u001b[0m \u001b[31m4.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m149.7/149.7 kB\u001b[0m \u001b[31m13.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m64.0/64.0 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m110.5/110.5 kB\u001b[0m \u001b[31m11.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m54.4/54.4 kB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.1/3.1 MB\u001b[0m \u001b[31m46.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m63.7/63.7 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m442.1/442.1 kB\u001b[0m \u001b[31m36.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.5/49.5 kB\u001b[0m \u001b[31m4.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m73.3/73.3 kB\u001b[0m \u001b[31m7.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m104.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m425.7/425.7 kB\u001b[0m \u001b[31m33.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m164.1/164.1 kB\u001b[0m \u001b[31m15.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m4.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for pypika (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "!pip install langchain langchain-community tiktoken chromadb openai pypdf sentence-transformers --quiet\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://constitutioncenter.org/media/files/constitution.pdf"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vYzVYv-WeHU1",
        "outputId": "0ecb89d7-c1c3-4314-80a1-dcb33b1d6593"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2024-11-04 01:23:03--  https://constitutioncenter.org/media/files/constitution.pdf\n",
            "Resolving constitutioncenter.org (constitutioncenter.org)... 104.22.22.181, 172.67.42.106, 104.22.23.181, ...\n",
            "Connecting to constitutioncenter.org (constitutioncenter.org)|104.22.22.181|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 413949 (404K) [application/pdf]\n",
            "Saving to: ‘constitution.pdf’\n",
            "\n",
            "\rconstitution.pdf      0%[                    ]       0  --.-KB/s               \rconstitution.pdf    100%[===================>] 404.25K  --.-KB/s    in 0.008s  \n",
            "\n",
            "2024-11-04 01:23:03 (50.7 MB/s) - ‘constitution.pdf’ saved [413949/413949]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Reading Legal Documents"
      ],
      "metadata": {
        "id": "WYrFsqTCyNci"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.document_loaders import PyPDFLoader\n",
        "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
        "from langchain.embeddings import OpenAIEmbeddings\n",
        "from langchain.vectorstores import Chroma\n",
        "from langchain.chains import RetrievalQA\n",
        "from langchain.llms import OpenAI\n",
        "import os\n",
        "\n",
        "# Set your OpenAI API key\n",
        "os.environ[\"OPENAI_API_KEY\"] = '<YOUR API KEY HERE>'\n",
        "\n",
        "# Step 1: Load the PDF and split it into manageable chunks\n",
        "def load_and_split_pdf(pdf_file_path):\n",
        "    loader = PyPDFLoader(pdf_file_path)\n",
        "    documents = loader.load()\n",
        "\n",
        "    # Split the document into smaller chunks for better processing\n",
        "    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)\n",
        "    texts = text_splitter.split_documents(documents)\n",
        "    return texts\n",
        "\n",
        "# Step 2: Store embeddings in ChromaDB\n",
        "def store_embeddings_in_chromadb(documents):\n",
        "    embeddings = OpenAIEmbeddings()  # You can also use SentenceTransformer embeddings if preferred\n",
        "    vectorstore = Chroma.from_documents(documents, embeddings)\n",
        "    return vectorstore\n",
        "\n",
        "# Step 3: Setup a retriever and chain for querying and context retrieval\n",
        "def create_retrieval_chain(vectorstore):\n",
        "    llm = OpenAI(model=\"gpt-3.5-turbo-instruct\",openai_api_key=os.environ[\"OPENAI_API_KEY\"])\n",
        "    retrieval_qa = RetrievalQA.from_chain_type(\n",
        "        llm=llm,\n",
        "        chain_type=\"stuff\",  # Optionally use \"map_reduce\" or others for more advanced chaining\n",
        "        retriever=vectorstore.as_retriever()\n",
        "    )\n",
        "    return retrieval_qa\n",
        "\n",
        "# Main Program\n",
        "if __name__ == \"__main__\":\n",
        "    # Sample PDF file path\n",
        "    pdf_file_path = 'constitution.pdf'\n",
        "\n",
        "    # Step 1: Load and split the PDF into chunks\n",
        "    print(\"Loading and splitting PDF...\")\n",
        "    documents = load_and_split_pdf(pdf_file_path)\n",
        "\n",
        "    # Step 2: Store document embeddings in ChromaDB\n",
        "    print(\"Storing embeddings in ChromaDB...\")\n",
        "    vectorstore = store_embeddings_in_chromadb(documents)\n",
        "\n",
        "    # Step 3: Create a retrieval-based chain with LLM\n",
        "    retrieval_qa = create_retrieval_chain(vectorstore)\n",
        "\n",
        "    # Step 4: User query input\n",
        "    user_query = input(\"Enter your query: \")\n",
        "\n",
        "    # Step 5: Retrieve context and generate answer using LLM\n",
        "    answer = retrieval_qa.run(user_query)\n",
        "    print(f\"LLM Response: {answer}\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tDtUcRwSc4KN",
        "outputId": "121429e2-ac5a-409c-bf6c-e8893c146444"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading and splitting PDF...\n",
            "Storing embeddings in ChromaDB...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-6-a55ecc9f15fa>:24: LangChainDeprecationWarning: The class `OpenAIEmbeddings` was deprecated in LangChain 0.0.9 and will be removed in 1.0. An updated version of the class exists in the :class:`~langchain-openai package and should be used instead. To use it run `pip install -U :class:`~langchain-openai` and import as `from :class:`~langchain_openai import OpenAIEmbeddings``.\n",
            "  embeddings = OpenAIEmbeddings()  # You can also use SentenceTransformer embeddings if preferred\n",
            "<ipython-input-6-a55ecc9f15fa>:30: LangChainDeprecationWarning: The class `OpenAI` was deprecated in LangChain 0.0.10 and will be removed in 1.0. An updated version of the class exists in the :class:`~langchain-openai package and should be used instead. To use it run `pip install -U :class:`~langchain-openai` and import as `from :class:`~langchain_openai import OpenAI``.\n",
            "  llm = OpenAI(model=\"gpt-3.5-turbo-instruct\",openai_api_key=os.environ[\"OPENAI_API_KEY\"])\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your query: WHat is the second ammendment?\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-6-a55ecc9f15fa>:58: LangChainDeprecationWarning: The method `Chain.run` was deprecated in langchain 0.1.0 and will be removed in 1.0. Use :meth:`~invoke` instead.\n",
            "  answer = retrieval_qa.run(user_query)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "LLM Response:  The second amendment to the Constitution guarantees the right of citizens to keep and bear arms.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Financial Reports Analyzer"
      ],
      "metadata": {
        "id": "Ei-qRfOdyGuy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModel\n",
        "from sentence_transformers import SentenceTransformer\n",
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import torch\n",
        "import re\n",
        "from typing import List, Dict, Tuple, Optional, Set\n",
        "from pathlib import Path\n",
        "import logging\n",
        "from datetime import datetime\n",
        "from dataclasses import dataclass\n",
        "from decimal import Decimal, InvalidOperation\n",
        "\n",
        "# Set up logging\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "@dataclass\n",
        "class FinancialMetric:\n",
        "    \"\"\"Data class to store financial metric information\"\"\"\n",
        "    name: str\n",
        "    value: float\n",
        "    unit: str\n",
        "    period: str\n",
        "    context: str\n",
        "\n",
        "class FinancialMetricsExtractor:\n",
        "    \"\"\"Class for extracting and calculating financial metrics\"\"\"\n",
        "\n",
        "    # Define common financial metrics patterns\n",
        "    METRIC_PATTERNS = {\n",
        "        'revenue': r'(?:revenue|sales)(?:\\s+of)?\\s+\\$?([\\d,]+(?:\\.\\d+)?)\\s*(?:billion|million|B|M)?',\n",
        "        'profit_margin': r'(?:profit\\s+margin|margin)\\s+(?:of\\s+)?(\\d+(?:\\.\\d+)?)\\%',\n",
        "        'operating_margin': r'operating\\s+margin\\s+(?:of\\s+)?(\\d+(?:\\.\\d+)?)\\%',\n",
        "        'net_income': r'(?:net\\s+income|profit)\\s+(?:of\\s+)?\\$?([\\d,]+(?:\\.\\d+)?)\\s*(?:billion|million|B|M)?',\n",
        "        'ebitda': r'(?:EBITDA|ebitda)\\s+(?:of\\s+)?\\$?([\\d,]+(?:\\.\\d+)?)\\s*(?:billion|million|B|M)?',\n",
        "        'eps': r'(?:EPS|earnings\\s+per\\s+share)\\s+(?:of\\s+)?\\$?([\\d,]+(?:\\.\\d+)?)',\n",
        "        'pe_ratio': r'(?:P/E|price[/-]to[/-]earnings)\\s+(?:ratio\\s+)?(?:of\\s+)?(\\d+(?:\\.\\d+)?)',\n",
        "        'debt_equity': r'debt[/-]to[/-]equity\\s+(?:ratio\\s+)?(?:of\\s+)?(\\d+(?:\\.\\d+)?)',\n",
        "        'current_ratio': r'current\\s+ratio\\s+(?:of\\s+)?(\\d+(?:\\.\\d+)?)',\n",
        "        'quick_ratio': r'quick\\s+ratio\\s+(?:of\\s+)?(\\d+(?:\\.\\d+)?)',\n",
        "        'roe': r'(?:ROE|return\\s+on\\s+equity)\\s+(?:of\\s+)?(\\d+(?:\\.\\d+)?)\\%',\n",
        "        'roa': r'(?:ROA|return\\s+on\\s+assets)\\s+(?:of\\s+)?(\\d+(?:\\.\\d+)?)\\%',\n",
        "        'cash_flow': r'(?:cash\\s+flow|operating\\s+cash\\s+flow)\\s+(?:of\\s+)?\\$?([\\d,]+(?:\\.\\d+)?)\\s*(?:billion|million|B|M)?',\n",
        "    }\n",
        "\n",
        "    # Unit multipliers\n",
        "    UNIT_MULTIPLIERS = {\n",
        "        'billion': 1e9,\n",
        "        'B': 1e9,\n",
        "        'million': 1e6,\n",
        "        'M': 1e6\n",
        "    }\n",
        "\n",
        "    @staticmethod\n",
        "    def normalize_number(value: str, unit: str = '') -> float:\n",
        "        \"\"\"\n",
        "        Normalize numerical values to a standard format\n",
        "\n",
        "        Args:\n",
        "            value (str): The numerical value as string\n",
        "            unit (str): The unit of the value (billion, million, etc.)\n",
        "\n",
        "        Returns:\n",
        "            float: Normalized value\n",
        "        \"\"\"\n",
        "        try:\n",
        "            # Remove commas and convert to float\n",
        "            clean_value = float(value.replace(',', ''))\n",
        "\n",
        "            # Apply unit multiplier if present\n",
        "            multiplier = FinancialMetricsExtractor.UNIT_MULTIPLIERS.get(unit, 1)\n",
        "            return clean_value * multiplier\n",
        "        except (ValueError, TypeError) as e:\n",
        "            logger.error(f\"Error normalizing number {value}: {str(e)}\")\n",
        "            return 0.0\n",
        "\n",
        "    @staticmethod\n",
        "    def extract_metrics(text: str) -> List[FinancialMetric]:\n",
        "        \"\"\"\n",
        "        Extract financial metrics from text\n",
        "\n",
        "        Args:\n",
        "            text (str): Text to analyze\n",
        "\n",
        "        Returns:\n",
        "            List[FinancialMetric]: List of extracted financial metrics\n",
        "        \"\"\"\n",
        "        metrics = []\n",
        "\n",
        "        # Extract year or period information\n",
        "        period_match = re.search(r'(?:fiscal\\s+year|FY|Q[1-4])?\\s*(?:20\\d{2})', text)\n",
        "        period = period_match.group(0) if period_match else \"Unknown\"\n",
        "\n",
        "        for metric_name, pattern in FinancialMetricsExtractor.METRIC_PATTERNS.items():\n",
        "            matches = re.finditer(pattern, text, re.IGNORECASE)\n",
        "\n",
        "            for match in matches:\n",
        "                value_str = match.group(1)\n",
        "\n",
        "                # Determine unit\n",
        "                unit = ''\n",
        "                for unit_name in FinancialMetricsExtractor.UNIT_MULTIPLIERS.keys():\n",
        "                    if unit_name in text[match.start():match.end()]:\n",
        "                        unit = unit_name\n",
        "                        break\n",
        "\n",
        "                # Normalize value\n",
        "                value = FinancialMetricsExtractor.normalize_number(value_str, unit)\n",
        "\n",
        "                # Get context (surrounding text)\n",
        "                start_idx = max(0, match.start() - 50)\n",
        "                end_idx = min(len(text), match.end() + 50)\n",
        "                context = text[start_idx:end_idx].strip()\n",
        "\n",
        "                metrics.append(FinancialMetric(\n",
        "                    name=metric_name,\n",
        "                    value=value,\n",
        "                    unit=unit,\n",
        "                    period=period,\n",
        "                    context=context\n",
        "                ))\n",
        "\n",
        "        return metrics\n",
        "\n",
        "class FinancialRatioCalculator:\n",
        "    \"\"\"Class for calculating financial ratios from extracted metrics\"\"\"\n",
        "\n",
        "    @staticmethod\n",
        "    def calculate_ratios(metrics: List[FinancialMetric]) -> Dict[str, float]:\n",
        "        \"\"\"\n",
        "        Calculate financial ratios from extracted metrics\n",
        "\n",
        "        Args:\n",
        "            metrics (List[FinancialMetric]): List of extracted metrics\n",
        "\n",
        "        Returns:\n",
        "            Dict[str, float]: Calculated ratios\n",
        "        \"\"\"\n",
        "        # Convert metrics to dictionary for easier access\n",
        "        metric_dict = {metric.name: metric.value for metric in metrics}\n",
        "        ratios = {}\n",
        "\n",
        "        # Gross Margin\n",
        "        if all(key in metric_dict for key in ['revenue', 'cost_of_goods']):\n",
        "            ratios['gross_margin'] = (metric_dict['revenue'] - metric_dict['cost_of_goods']) / metric_dict['revenue']\n",
        "\n",
        "        # Operating Margin\n",
        "        if all(key in metric_dict for key in ['operating_income', 'revenue']):\n",
        "            ratios['operating_margin'] = metric_dict['operating_income'] / metric_dict['revenue']\n",
        "\n",
        "        # Return on Equity (ROE)\n",
        "        if all(key in metric_dict for key in ['net_income', 'shareholders_equity']):\n",
        "            ratios['roe'] = metric_dict['net_income'] / metric_dict['shareholders_equity']\n",
        "\n",
        "        return ratios\n",
        "\n",
        "\n",
        "class FinancialReportAnalyzer:\n",
        "    def __init__(self, model_name: str = \"sentence-transformers/all-MiniLM-L6-v2\"):\n",
        "        try:\n",
        "            self.embedding_model = SentenceTransformer(model_name)\n",
        "            self.document_chunks: List[str] = []\n",
        "            self.embeddings: Optional[np.ndarray] = None\n",
        "            self.metadata: Dict = {}\n",
        "            self.metrics_extractor = FinancialMetricsExtractor()\n",
        "            self.ratio_calculator = FinancialRatioCalculator()\n",
        "            logger.info(f\"Initialized FinancialReportAnalyzer with model: {model_name}\")\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Failed to initialize model: {str(e)}\")\n",
        "            raise RuntimeError(f\"Model initialization failed: {str(e)}\")\n",
        "\n",
        "    def preprocess_document(self, text: str, chunk_size: int = 500) -> List[str]:\n",
        "        \"\"\"\n",
        "        Preprocess and chunk the financial report text.\n",
        "\n",
        "        Args:\n",
        "            text (str): Raw text from financial report\n",
        "            chunk_size (int): Maximum size of each text chunk\n",
        "\n",
        "        Returns:\n",
        "            List[str]: List of preprocessed text chunks\n",
        "\n",
        "        Raises:\n",
        "            ValueError: If text is empty or chunk_size is invalid\n",
        "        \"\"\"\n",
        "        if not text or not isinstance(text, str):\n",
        "            raise ValueError(\"Input text must be a non-empty string\")\n",
        "        if chunk_size <= 0:\n",
        "            raise ValueError(\"Chunk size must be positive\")\n",
        "\n",
        "        # Remove extra whitespace and normalize text\n",
        "        text = re.sub(r'\\s+', ' ', text)\n",
        "\n",
        "        # Preserve important financial characters and formatting\n",
        "        text = re.sub(r'[^\\w\\s.,;:()$%+-]', '', text)\n",
        "\n",
        "        # Improved sentence splitting that preserves decimal numbers and currency\n",
        "        sentences = re.split(r'(?<!\\d)\\.(?!\\d)(?!\\s*[a-z])', text)\n",
        "        chunks = []\n",
        "        current_chunk = \"\"\n",
        "\n",
        "        # Enhanced chunking logic that preserves context and financial information\n",
        "        for sentence in sentences:\n",
        "            sentence = sentence.strip()\n",
        "            if not sentence:\n",
        "                continue\n",
        "\n",
        "            # Check if adding this sentence would exceed chunk size\n",
        "            if len(current_chunk) + len(sentence) < chunk_size:\n",
        "                current_chunk += sentence + \". \"\n",
        "            else:\n",
        "                if current_chunk:\n",
        "                    chunks.append(current_chunk.strip())\n",
        "                current_chunk = sentence + \". \"\n",
        "\n",
        "        if current_chunk:\n",
        "            chunks.append(current_chunk.strip())\n",
        "\n",
        "        # Validate chunks\n",
        "        if not chunks:\n",
        "            raise ValueError(\"No valid chunks were generated from the input text\")\n",
        "\n",
        "        logger.info(f\"Created {len(chunks)} chunks from document\")\n",
        "        return chunks\n",
        "\n",
        "    def create_embeddings(self, chunks: List[str]) -> None:\n",
        "        \"\"\"\n",
        "        Create embeddings for all document chunks.\n",
        "\n",
        "        Args:\n",
        "            chunks (List[str]): List of text chunks to embed\n",
        "\n",
        "        Raises:\n",
        "            ValueError: If chunks list is empty\n",
        "        \"\"\"\n",
        "        if not chunks:\n",
        "            raise ValueError(\"Cannot create embeddings from empty chunks list\")\n",
        "\n",
        "        try:\n",
        "            self.document_chunks = chunks\n",
        "            self.embeddings = self.embedding_model.encode(chunks, show_progress_bar=True)\n",
        "            logger.info(f\"Created embeddings of shape: {self.embeddings.shape}\")\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Failed to create embeddings: {str(e)}\")\n",
        "            raise\n",
        "\n",
        "    def retrieve_relevant_chunks(self, query: str, n_chunks: int = 3,\n",
        "                               similarity_threshold: float = 0.3) -> List[Tuple[str, float]]:\n",
        "        \"\"\"\n",
        "        Retrieve the most relevant chunks for a given query.\n",
        "\n",
        "        Args:\n",
        "            query (str): User query about financial information\n",
        "            n_chunks (int): Number of relevant chunks to retrieve\n",
        "            similarity_threshold (float): Minimum similarity score to include chunk\n",
        "\n",
        "        Returns:\n",
        "            List[Tuple[str, float]]: List of (chunk, similarity_score) tuples\n",
        "        \"\"\"\n",
        "        if not query.strip():\n",
        "            raise ValueError(\"Query cannot be empty\")\n",
        "        if self.embeddings is None:\n",
        "            raise ValueError(\"No embeddings available. Please load a document first.\")\n",
        "\n",
        "        try:\n",
        "            query_embedding = self.embedding_model.encode([query])[0]\n",
        "            similarities = cosine_similarity([query_embedding], self.embeddings)[0]\n",
        "\n",
        "            # Filter by similarity threshold and get top chunks\n",
        "            valid_indices = np.where(similarities >= similarity_threshold)[0]\n",
        "            top_indices = valid_indices[np.argsort(similarities[valid_indices])[-n_chunks:][::-1]]\n",
        "\n",
        "            results = [(self.document_chunks[i], similarities[i]) for i in top_indices]\n",
        "            logger.info(f\"Retrieved {len(results)} relevant chunks for query\")\n",
        "            return results\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Error retrieving chunks: {str(e)}\")\n",
        "            raise\n",
        "\n",
        "    def analyze_query(self, query: str) -> Dict:\n",
        "        \"\"\"\n",
        "        Analyze a specific query about the financial report.\n",
        "\n",
        "        Args:\n",
        "            query (str): User query about financial information\n",
        "\n",
        "        Returns:\n",
        "            Dict: Analysis results including relevant chunks, metrics, and ratios\n",
        "        \"\"\"\n",
        "        if not query.strip():\n",
        "            raise ValueError(\"Query cannot be empty\")\n",
        "\n",
        "        try:\n",
        "            relevant_chunks = self.retrieve_relevant_chunks(query)\n",
        "\n",
        "            # Extract metrics from relevant chunks\n",
        "            all_metrics = []\n",
        "            for chunk, _ in relevant_chunks:\n",
        "                metrics = FinancialMetricsExtractor.extract_metrics(chunk)\n",
        "                all_metrics.extend(metrics)\n",
        "\n",
        "            # Calculate relevant ratios\n",
        "            ratios = FinancialRatioCalculator.calculate_ratios(all_metrics)\n",
        "\n",
        "            # Organize metrics by type\n",
        "            metrics_by_type = {}\n",
        "            for metric in all_metrics:\n",
        "                if metric.name not in metrics_by_type:\n",
        "                    metrics_by_type[metric.name] = []\n",
        "                metrics_by_type[metric.name].append({\n",
        "                    'value': metric.value,\n",
        "                    'unit': metric.unit,\n",
        "                    'period': metric.period,\n",
        "                    'context': metric.context\n",
        "                })\n",
        "\n",
        "            analysis_result = {\n",
        "                \"query\": query,\n",
        "                \"relevant_sections\": [\n",
        "                    {\n",
        "                        \"text\": chunk,\n",
        "                        \"relevance_score\": float(score),\n",
        "                    }\n",
        "                    for chunk, score in relevant_chunks\n",
        "                ],\n",
        "                \"extracted_metrics\": metrics_by_type,\n",
        "                \"calculated_ratios\": ratios,\n",
        "                \"metadata\": self.metadata\n",
        "            }\n",
        "\n",
        "            return analysis_result\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Error analyzing query: {str(e)}\")\n",
        "            raise\n",
        "\n",
        "\n",
        "\n",
        "def main():\n",
        "    \"\"\"Example usage with financial metrics extraction\"\"\"\n",
        "    try:\n",
        "        analyzer = FinancialReportAnalyzer()\n",
        "\n",
        "        sample_report = \"\"\"\n",
        "        In fiscal year 2023, our company achieved strong financial performance with revenue growth of 15% year-over-year,\n",
        "        reaching $2.5 billion. Our operating margin improved to 28%, up from 25% in the previous year.\n",
        "        The company's cash flow from operations was $500 million, representing a 20% increase from 2022.\n",
        "\n",
        "        Our EBITDA reached $750 million with a margin of 30%. The P/E ratio stands at 22.5, reflecting strong market confidence.\n",
        "        Net income increased to $400 million, resulting in earnings per share (EPS) of $2.45.\n",
        "\n",
        "        Return on equity (ROE) improved to 18.5% while return on assets (ROA) reached 12.3%.\n",
        "        Our debt-to-equity ratio decreased to 0.8, showing improved financial health.\n",
        "        \"\"\"\n",
        "\n",
        "        chunks = analyzer.preprocess_document(sample_report)\n",
        "        analyzer.create_embeddings(chunks)\n",
        "\n",
        "        # Example queries focusing on financial metrics\n",
        "        queries = [\n",
        "            \"What are the company's profitability metrics?\",\n",
        "            \"What are the key financial ratios?\",\n",
        "            \"How is the company's operational performance?\"\n",
        "        ]\n",
        "\n",
        "        for query in queries:\n",
        "            result = analyzer.analyze_query(query)\n",
        "            print(f\"\\nQuery: {result['query']}\")\n",
        "            print(\"\\nExtracted Metrics:\")\n",
        "            for metric_type, metrics in result['extracted_metrics'].items():\n",
        "                print(f\"\\n{metric_type.upper()}:\")\n",
        "                for metric in metrics:\n",
        "                    print(f\"- Value: {metric['value']}\")\n",
        "                    print(f\"  Period: {metric['period']}\")\n",
        "                    print(f\"  Context: {metric['context']}\")\n",
        "\n",
        "            if result['calculated_ratios']:\n",
        "                print(\"\\nCalculated Ratios:\")\n",
        "                for ratio_name, value in result['calculated_ratios'].items():\n",
        "                    print(f\"- {ratio_name}: {value:.2f}\")\n",
        "\n",
        "    except Exception as e:\n",
        "        logger.error(f\"Error in main: {str(e)}\")\n",
        "        raise\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 527,
          "referenced_widgets": [
            "1d1fbb9e6c6649188b3eb8a7ff64ea40",
            "3b7017e547ae4e418b6296f6421d8a36",
            "62d29109e30148cd8d33b9d0b58d5b81",
            "30694df5cbdb4e00b70534f963ec61fd",
            "11ff540c531c48f08b241beb8e5e3a19",
            "36407cf7a73d466abbe732b7681a4de1",
            "7bd1f4447764435cbbdcde2d587906cb",
            "a517d1a8c2de4333850843670a8505ba",
            "f4064c8c509e4bc0887a2222e9b2c3e2",
            "7ce67680963c4c1fafc2df67bfe97710",
            "e708fe1190174d7597082288d6a9cc1a"
          ]
        },
        "id": "wUrH_AEKvvYT",
        "outputId": "ae6a5515-622a-4f1c-8240-6e06ec1f094a"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "1d1fbb9e6c6649188b3eb8a7ff64ea40"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Query: What are the company's profitability metrics?\n",
            "\n",
            "Extracted Metrics:\n",
            "\n",
            "PROFIT_MARGIN:\n",
            "- Value: 30.0\n",
            "  Period: fiscal year 2023\n",
            "  Context: from 2022. Our EBITDA reached $750 million with a margin of 30%. The PE ratio stands at 22.5, reflecting strong m\n",
            "\n",
            "Query: What are the key financial ratios?\n",
            "\n",
            "Extracted Metrics:\n",
            "\n",
            "PROFIT_MARGIN:\n",
            "- Value: 30.0\n",
            "  Period: fiscal year 2023\n",
            "  Context: from 2022. Our EBITDA reached $750 million with a margin of 30%. The PE ratio stands at 22.5, reflecting strong m\n",
            "\n",
            "Query: How is the company's operational performance?\n",
            "\n",
            "Extracted Metrics:\n",
            "\n",
            "PROFIT_MARGIN:\n",
            "- Value: 30.0\n",
            "  Period: fiscal year 2023\n",
            "  Context: from 2022. Our EBITDA reached $750 million with a margin of 30%. The PE ratio stands at 22.5, reflecting strong m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Scientific Document Summary"
      ],
      "metadata": {
        "id": "9vM4tpTC3bma"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Dict, List, Optional, Tuple\n",
        "import numpy as np\n",
        "from sentence_transformers import SentenceTransformer\n",
        "from transformers import pipeline\n",
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "import re\n",
        "from dataclasses import dataclass\n",
        "import torch\n",
        "import logging\n",
        "\n",
        "logging.basicConfig(level=logging.INFO)\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "@dataclass\n",
        "class RelevantPassage:\n",
        "    section: str\n",
        "    text: str\n",
        "    similarity: float\n",
        "\n",
        "class ResearchPaperRAG:\n",
        "    def __init__(self,\n",
        "                 embedding_model: str = 'all-MiniLM-L6-v2',\n",
        "                 summarizer_model: str = \"facebook/bart-large-cnn\"):\n",
        "        \"\"\"Initialize the RAG system with specified models.\"\"\"\n",
        "        try:\n",
        "            # Initialize models without explicit device setting first\n",
        "            self.embedding_model = SentenceTransformer(embedding_model)\n",
        "\n",
        "            # Initialize summarizer pipeline\n",
        "            self.summarizer = pipeline(\n",
        "                \"summarization\",\n",
        "                model=summarizer_model\n",
        "            )\n",
        "\n",
        "            logger.info(\"Models initialized successfully\")\n",
        "\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Model initialization failed: {str(e)}\")\n",
        "            raise RuntimeError(f\"Error initializing models: {str(e)}\")\n",
        "\n",
        "    def get_embeddings(self, text_chunks: List[str]) -> np.ndarray:\n",
        "        \"\"\"Generate embeddings for text chunks.\"\"\"\n",
        "        if not text_chunks:\n",
        "            raise ValueError(\"Text chunks list cannot be empty\")\n",
        "\n",
        "        try:\n",
        "            return self.embedding_model.encode(\n",
        "                text_chunks,\n",
        "                show_progress_bar=False,\n",
        "                convert_to_numpy=True\n",
        "            )\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Embedding generation failed: {str(e)}\")\n",
        "            raise RuntimeError(f\"Error generating embeddings: {str(e)}\")\n",
        "\n",
        "    def preprocess_paper(self, paper_text: str) -> Dict[str, str]:\n",
        "        \"\"\"Extract sections from the research paper.\"\"\"\n",
        "        if not paper_text or not isinstance(paper_text, str):\n",
        "            raise ValueError(\"Paper text must be a non-empty string\")\n",
        "\n",
        "        # Define section patterns with variations\n",
        "        section_patterns = {\n",
        "            'abstract': r'abstract|summary',\n",
        "            'introduction': r'introduction|background',\n",
        "            'methodology': r'methodology|methods|materials and methods|experimental procedure',\n",
        "            'results': r'results|findings',\n",
        "            'discussion': r'discussion|interpretation',\n",
        "            'conclusion': r'conclusion|conclusions|final remarks',\n",
        "            'limitations': r'limitations|study limitations|constraints'\n",
        "        }\n",
        "\n",
        "        sections = {k: '' for k in section_patterns.keys()}\n",
        "        current_section = None\n",
        "        current_text = []\n",
        "\n",
        "        lines = paper_text.split('\\n')\n",
        "        for line in lines:\n",
        "            line = line.strip()\n",
        "            if not line:\n",
        "                continue\n",
        "\n",
        "            # Check if line is a section header\n",
        "            found_section = False\n",
        "            for section, pattern in section_patterns.items():\n",
        "                if re.match(f\"^(?i)({pattern})\\s*:?$\", line):\n",
        "                    if current_section and current_text:\n",
        "                        sections[current_section] += ' '.join(current_text)\n",
        "                    current_section = section\n",
        "                    current_text = []\n",
        "                    found_section = True\n",
        "                    break\n",
        "\n",
        "            if not found_section and current_section:\n",
        "                current_text.append(line)\n",
        "\n",
        "        # Add the last section\n",
        "        if current_section and current_text:\n",
        "            sections[current_section] += ' '.join(current_text)\n",
        "\n",
        "        # Clean up empty sections and normalize whitespace\n",
        "        return {k: ' '.join(v.split()) for k, v in sections.items() if v.strip()}\n",
        "\n",
        "    def chunk_text(self, text: str, chunk_size: int = 1000) -> List[str]:\n",
        "        \"\"\"Split text into chunks while preserving sentence boundaries.\"\"\"\n",
        "        if not text:\n",
        "            return []\n",
        "\n",
        "        sentences = re.split(r'(?<=[.!?])\\s+(?=[A-Z])|(?<=\\n)\\s*(?=[A-Z])', text)\n",
        "        chunks = []\n",
        "        current_chunk = []\n",
        "        current_size = 0\n",
        "\n",
        "        for sentence in sentences:\n",
        "            sentence = sentence.strip()\n",
        "            if not sentence:\n",
        "                continue\n",
        "\n",
        "            sentence_size = len(sentence)\n",
        "\n",
        "            if sentence_size > chunk_size:\n",
        "                if current_chunk:\n",
        "                    chunks.append(' '.join(current_chunk))\n",
        "                    current_chunk = []\n",
        "                    current_size = 0\n",
        "\n",
        "                clauses = re.split(r'(?<=[,;:])\\s+', sentence)\n",
        "                for clause in clauses:\n",
        "                    if len(clause) > chunk_size:\n",
        "                        chunks.append(clause)\n",
        "                    elif current_size + len(clause) > chunk_size:\n",
        "                        chunks.append(' '.join(current_chunk))\n",
        "                        current_chunk = [clause]\n",
        "                        current_size = len(clause)\n",
        "                    else:\n",
        "                        current_chunk.append(clause)\n",
        "                        current_size += len(clause)\n",
        "                continue\n",
        "\n",
        "            if current_size + sentence_size > chunk_size and current_chunk:\n",
        "                chunks.append(' '.join(current_chunk))\n",
        "                current_chunk = [sentence]\n",
        "                current_size = sentence_size\n",
        "            else:\n",
        "                current_chunk.append(sentence)\n",
        "                current_size += sentence_size\n",
        "\n",
        "        if current_chunk:\n",
        "            chunks.append(' '.join(current_chunk))\n",
        "\n",
        "        return [chunk for chunk in chunks if chunk.strip()]\n",
        "\n",
        "    def retrieve_relevant_passages(self,\n",
        "                                 query: str,\n",
        "                                 sections: Dict[str, str],\n",
        "                                 top_k: int = 3,\n",
        "                                 similarity_threshold: float = 0.2) -> List[RelevantPassage]:\n",
        "        \"\"\"Retrieve most relevant passages based on query.\"\"\"\n",
        "        if not query or not sections:\n",
        "            raise ValueError(\"Query and sections cannot be empty\")\n",
        "\n",
        "        relevant_passages = []\n",
        "        query_embedding = self.embedding_model.encode([query])[0]\n",
        "\n",
        "        for section_name, section_text in sections.items():\n",
        "            chunks = self.chunk_text(section_text)\n",
        "            if not chunks:\n",
        "                continue\n",
        "\n",
        "            chunk_embeddings = self.get_embeddings(chunks)\n",
        "            similarities = cosine_similarity([query_embedding], chunk_embeddings)[0]\n",
        "\n",
        "            for idx, chunk in enumerate(chunks):\n",
        "                if similarities[idx] > similarity_threshold:\n",
        "                    relevant_passages.append(RelevantPassage(\n",
        "                        section=section_name,\n",
        "                        text=chunk,\n",
        "                        similarity=float(similarities[idx])\n",
        "                    ))\n",
        "\n",
        "        # Sort by similarity and return top_k\n",
        "        relevant_passages.sort(key=lambda x: x.similarity, reverse=True)\n",
        "        return relevant_passages[:top_k]\n",
        "\n",
        "    def generate_summary(self, passages: List[RelevantPassage], query: Optional[str] = None) -> str:\n",
        "        \"\"\"Generate a summary based on retrieved passages.\"\"\"\n",
        "        if not passages:\n",
        "            return \"No relevant passages found to summarize.\"\n",
        "\n",
        "        try:\n",
        "            combined_text = \" \".join([p.text for p in passages])\n",
        "            input_length = len(combined_text.split())\n",
        "\n",
        "            max_length = min(150, max(50, input_length // 2))\n",
        "            min_length = max(25, input_length // 4)\n",
        "\n",
        "            summary = self.summarizer(\n",
        "                combined_text,\n",
        "                max_length=max_length,\n",
        "                min_length=min_length,\n",
        "                do_sample=False,\n",
        "                truncation=True\n",
        "            )[0]['summary_text']\n",
        "\n",
        "            return summary\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Error in summary generation: {str(e)}\")\n",
        "            return \"Error generating summary. Showing most relevant passage instead: \" + passages[0].text\n",
        "\n",
        "    def query_paper(self, paper_text: str, query: str, top_k: int = 3) -> Dict[str, any]:\n",
        "        \"\"\"Main method to query a paper.\"\"\"\n",
        "        if not paper_text or not query:\n",
        "            raise ValueError(\"Paper text and query must be non-empty strings\")\n",
        "\n",
        "        try:\n",
        "            sections = self.preprocess_paper(paper_text)\n",
        "\n",
        "            if not sections:\n",
        "                raise ValueError(\"No valid sections found in paper\")\n",
        "\n",
        "            relevant_passages = self.retrieve_relevant_passages(\n",
        "                query=query,\n",
        "                sections=sections,\n",
        "                top_k=top_k\n",
        "            )\n",
        "\n",
        "            summary = self.generate_summary(relevant_passages, query)\n",
        "\n",
        "            return {\n",
        "                'summary': summary,\n",
        "                'relevant_passages': relevant_passages\n",
        "            }\n",
        "\n",
        "        except Exception as e:\n",
        "            logger.error(f\"Error processing paper: {str(e)}\")\n",
        "            raise RuntimeError(f\"Error processing paper: {str(e)}\")\n",
        "\n",
        "def main():\n",
        "    \"\"\"Example usage of the ResearchPaperRAG system.\"\"\"\n",
        "    try:\n",
        "        # Sample paper with more detailed methodology section\n",
        "        sample_paper = \"\"\"\n",
        "        Abstract:\n",
        "        This comprehensive study investigates the effects of caffeine on cognitive performance and reaction time in healthy adults.\n",
        "\n",
        "        Introduction:\n",
        "        Caffeine is one of the most widely consumed psychoactive substances globally.\n",
        "\n",
        "        Methodology:\n",
        "        We conducted a double-blind, placebo-controlled study with 100 participants aged 18-65. Participants were randomly assigned to either the caffeine group (200mg) or placebo group. We used a computerized cognitive assessment battery to measure performance. Testing sessions occurred at baseline and 1 hour post-administration. Each participant completed three testing blocks: memory tasks, reaction time assessment, and attention span measurement.\n",
        "\n",
        "        Results:\n",
        "        The results showed significant improvement in reaction times (p<0.001).\n",
        "\n",
        "        Discussion:\n",
        "        Our findings suggest that moderate caffeine consumption enhances cognitive performance.\n",
        "\n",
        "        Limitations:\n",
        "        The study was limited by its single-dose design and short duration.\n",
        "\n",
        "        Conclusion:\n",
        "        This study demonstrates the positive effects of caffeine on cognitive performance.\n",
        "        \"\"\"\n",
        "\n",
        "        rag = ResearchPaperRAG()\n",
        "\n",
        "        queries = [\n",
        "            \"What were the main findings about caffeine's effects?\",\n",
        "            \"Describe the methodology used in the study\",\n",
        "            \"What were the limitations of the study?\"\n",
        "        ]\n",
        "\n",
        "        for query in queries:\n",
        "            print(f\"\\nQuery: {query}\")\n",
        "            result = rag.query_paper(sample_paper, query)\n",
        "            print(\"\\nSummary:\", result['summary'])\n",
        "            print(\"\\nRelevant Passages:\")\n",
        "            for passage in result['relevant_passages']:\n",
        "                print(f\"\\nFrom {passage.section} section (similarity: {passage.similarity:.2f}):\")\n",
        "                print(passage.text)\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Error in main: {str(e)}\")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fSmSnhbhyWcm",
        "outputId": "32b26ce2-f629-465b-ba20-13a129b87adb"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Hardware accelerator e.g. GPU is available in the environment, but no `device` argument is passed to the `Pipeline` object. Model will be on CPU.\n",
            "Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n",
            "Your max_length is set to 50, but your input_length is only 47. Since this is a summarization task, where outputs shorter than the input are typically wanted, you might consider decreasing max_length manually, e.g. summarizer('...', max_length=23)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Query: What were the main findings about caffeine's effects?\n",
            "\n",
            "Summary: Caffeine is one of the most widely consumed psychoactive substances globally. This study demonstrates the positive effects of caffeine on cognitive performance.\n",
            "\n",
            "Relevant Passages:\n",
            "\n",
            "From conclusion section (similarity: 0.69):\n",
            "This study demonstrates the positive effects of caffeine on cognitive performance.\n",
            "\n",
            "From abstract section (similarity: 0.64):\n",
            "This comprehensive study investigates the effects of caffeine on cognitive performance and reaction time in healthy adults.\n",
            "\n",
            "From introduction section (similarity: 0.62):\n",
            "Caffeine is one of the most widely consumed psychoactive substances globally.\n",
            "\n",
            "Query: Describe the methodology used in the study\n",
            "\n",
            "Summary: We conducted a double-blind, placebo-controlled study with 100 participants aged 18-65. Participants were randomly assigned to either the caffeine group (200mg) or placebo group. We used a computerized cognitive assessment battery to measure performance\n",
            "\n",
            "Relevant Passages:\n",
            "\n",
            "From limitations section (similarity: 0.29):\n",
            "The study was limited by its single-dose design and short duration.\n",
            "\n",
            "From methodology section (similarity: 0.26):\n",
            "We conducted a double-blind, placebo-controlled study with 100 participants aged 18-65. Participants were randomly assigned to either the caffeine group (200mg) or placebo group. We used a computerized cognitive assessment battery to measure performance. Testing sessions occurred at baseline and 1 hour post-administration. Each participant completed three testing blocks: memory tasks, reaction time assessment, and attention span measurement.\n",
            "\n",
            "Query: What were the limitations of the study?\n",
            "\n",
            "Summary: We conducted a double-blind, placebo-controlled study with 100 participants aged 18-65. Participants were randomly assigned to either the caffeine group (200mg) or placebo group. We used a computerized cognitive assessment battery to measure performance\n",
            "\n",
            "Relevant Passages:\n",
            "\n",
            "From limitations section (similarity: 0.58):\n",
            "The study was limited by its single-dose design and short duration.\n",
            "\n",
            "From methodology section (similarity: 0.28):\n",
            "We conducted a double-blind, placebo-controlled study with 100 participants aged 18-65. Participants were randomly assigned to either the caffeine group (200mg) or placebo group. We used a computerized cognitive assessment battery to measure performance. Testing sessions occurred at baseline and 1 hour post-administration. Each participant completed three testing blocks: memory tasks, reaction time assessment, and attention span measurement.\n"
          ]
        }
      ]
    }
  ]
}